package org.liuyuantao.lucene.search;

import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.document.*;
import org.apache.lucene.index.*;
import org.apache.lucene.queryparser.classic.ParseException;
import org.apache.lucene.queryparser.classic.QueryParser;
import org.apache.lucene.search.BooleanClause.Occur;
import org.apache.lucene.search.*;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.util.BytesRef;

import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;

/**
 * 简单的查询API
 */
public class SearcherExample {
    public static final String INDEX_PATH = "src/main/resources/index/";
    private Directory directory;
    private IndexReader reader;
    private String[] ids = {"1", "2", "3", "4", "5", "6"};
    private String[] emails = {"aa@gmail.com", "bb@gmail.org", "cc@outlook.org", "dd@sina.org", "ee@163.edu", "ff@qq.org"};
    private String[] contents = {
            "welcome to visited the space,I am book",
            "hello boy, I like pingpeng ball",
            "my name is cc I like game",
            "I like football",
            "I like football and I like basketball too",
            "I like movie and swim"
    };
    private Date[] dates = null;
    private int[] attachs = {2, 3, 1, 4, 5, 5};
    private String[] names = {"zhangsan", "lisi", "sohn", "jetty", "mike", "jake"};
    private Map<String, Float> scores = new HashMap<String, Float>();

    /**
     * 实例化数据的时候初始化索引数据
     */
    public SearcherExample() {
        Path path = Paths.get(INDEX_PATH);
        try {
            directory = FSDirectory.open(path);
            scores.put("qq.org", 2.0f);
            scores.put("163.edu", 1.5f);
            setDates();
            index();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 初始化模拟日期数据
     */
    private void setDates() {
        SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd");
        try {
            dates = new Date[ids.length];
            dates[0] = sdf.parse("2017-02-19");
            dates[1] = sdf.parse("2017-01-11");
            dates[2] = sdf.parse("2017-09-19");
            dates[3] = sdf.parse("2016-12-22");
            dates[4] = sdf.parse("2015-01-01");
            dates[5] = sdf.parse("2014-05-19");
        } catch (java.text.ParseException e) {
            e.printStackTrace();
        }
    }

    /**
     * 创建索引
     */
    public void index() {
        IndexWriter writer = null;
        try {
            writer = new IndexWriter(directory, new IndexWriterConfig(new StandardAnalyzer()));
            writer.deleteAll();
            for (int i = 0; i < ids.length; i++) {
                Document doc = new Document();
                doc.add(new StringField("id", ids[i], Field.Store.YES));
                doc.add(new TextField("content", contents[i], Field.Store.NO));
                doc.add(new StringField("name", names[i], Field.Store.YES));
                //存储数字
                doc.add(new IntField("attach", attachs[i], Field.Store.YES));
                //存储日期
                doc.add(new LongField("date", dates[i].getTime(), Field.Store.YES));
                String et = emails[i].substring(emails[i].lastIndexOf("@") + 1);
                TextField emailField = new TextField("email", emails[i], Field.Store.YES);
                emailField.setBoost(scores.getOrDefault(et, 0.5f));
                doc.add(emailField);
                writer.addDocument(doc);
            }
            int max = writer.maxDoc();
            System.out.println("创建Document数:" + max);
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (writer != null) writer.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 根据默认的directory设置IndexReader
     *
     * @return
     */
    private IndexReader initIndexReader() {
        try {
            if (this.reader == null) {
                this.reader = DirectoryReader.open(directory);
            } else {
                IndexReader tr = DirectoryReader.openIfChanged((DirectoryReader) this.reader);
                if (tr != null) {
                    this.reader.close();
                    this.reader = tr;
                }
            }
            return this.reader;
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }


    /**
     * 字符串精确查找
     *
     * @param field
     * @param value
     * @param num
     */
    public void searchByTerm(String field, String value, int num) {
        try {
            initIndexReader();
            IndexSearcher searcher = new IndexSearcher(this.reader);
            Query query = new TermQuery(new Term(field, value));
            TopDocs tds = searcher.search(query, num);
            System.out.println("一共查询到数量:" + tds.totalHits);
            for (ScoreDoc sd : tds.scoreDocs) {
                Document doc = searcher.doc(sd.doc);
                System.out.println(doc.get("id") + "---->" +
                        doc.get("name") + "[" + doc.get("email") + "]");
            }
            this.reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 字符串区间查找
     *
     * @param field
     * @param start
     * @param end
     * @param num
     */
    public void searchByTermRange(String field, String start, String end, int num) {
        try {
            initIndexReader();
            IndexSearcher searcher = new IndexSearcher(this.reader);
            Query query = new TermRangeQuery(field, new BytesRef(start.getBytes()), new BytesRef(end.getBytes()), true, true);
            TopDocs tds = searcher.search(query, num);
            System.out.println("一共查询了:" + tds.totalHits);
            for (ScoreDoc sd : tds.scoreDocs) {
                Document doc = searcher.doc(sd.doc);
                System.out.println(doc.get("id") + "---->" +
                        doc.get("name") + "[" + doc.get("email") + "]");
            }
            this.reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 数字区间查找,只能用于数字区间查询
     *
     * @param field
     * @param start
     * @param end
     * @param num
     */
    public void searchByNumricRange(String field, int start, int end, int num) {
        try {
            initIndexReader();
            IndexSearcher searcher = new IndexSearcher(this.reader);
            Query query = NumericRangeQuery.newIntRange(field, start, end, true, true);
            TopDocs tds = searcher.search(query, num);
            System.out.println("一共查询了:" + tds.totalHits);
            for (ScoreDoc sd : tds.scoreDocs) {
                Document doc = searcher.doc(sd.doc);
                System.out.println(doc.get("id") + "---->" +
                        doc.get("name") + "[" + doc.get("email") + "]");
            }
            this.reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 前缀查找
     *
     * @param field
     * @param value
     * @param num
     */
    public void searchByPrefix(String field, String value, int num) {
        try {
            initIndexReader();
            IndexSearcher searcher = new IndexSearcher(this.reader);
            Query query = new PrefixQuery(new Term(field, value));
            TopDocs tds = searcher.search(query, num);
            System.out.println("一共查询了:" + tds.totalHits);
            for (ScoreDoc sd : tds.scoreDocs) {
                Document doc = searcher.doc(sd.doc);
                System.out.println(doc.get("id") + "---->" +
                        doc.get("name") + "[" + doc.get("email") + "]");
            }
            this.reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 通配符查找
     *
     * @param field
     * @param value
     * @param num
     */
    public void searchByWildcard(String field, String value, int num) {
        try {
            initIndexReader();
            IndexSearcher searcher = new IndexSearcher(this.reader);
            //在传入的value中可以使用通配符:?和*,?表示匹配一个字符，*表示匹配任意多个字符
            Query query = new WildcardQuery(new Term(field, value));
            TopDocs tds = searcher.search(query, num);
            System.out.println("一共查询了:" + tds.totalHits);
            for (ScoreDoc sd : tds.scoreDocs) {
                Document doc = searcher.doc(sd.doc);
                System.out.println(doc.get("id") + "---->" +
                        doc.get("name") + "[" + doc.get("email") + "]");
            }
            this.reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * Boolean查询，5.x版本使用Builder
     *
     * @param num
     */
    public void searchByBoolean(int num) {
        try {
            initIndexReader();
            IndexSearcher searcher = new IndexSearcher(this.reader);
            BooleanQuery.Builder builder = new BooleanQuery.Builder();
            /*
             * BooleanQuery可以连接多个子查询
			 * Occur.MUST表示必须出现
			 * Occur.SHOULD表示可以出现
			 * Occur.MUSE_NOT表示不能出现
			 */
            builder.add(new TermQuery(new Term("name", "zhangsan")), Occur.MUST_NOT);
            builder.add(new TermQuery(new Term("content", "game")), Occur.SHOULD);
            BooleanQuery query = builder.build();
            TopDocs tds = searcher.search(query, num);
            System.out.println("一共查询了:" + tds.totalHits);
            for (ScoreDoc sd : tds.scoreDocs) {
                Document doc = searcher.doc(sd.doc);
                System.out.println(doc.get("id") + "---->" +
                        doc.get("name") + "[" + doc.get("email") + "]-->" + doc.get("id") + "," +
                        doc.get("attach") + "," + doc.get("date"));
            }
            this.reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 单词查询,对中文支持不好
     *
     * @param num
     */
    public void searchByPhrase(int num) {
        try {
            initIndexReader();
            IndexSearcher searcher = new IndexSearcher(this.reader);
            PhraseQuery.Builder builder = new PhraseQuery.Builder();
            builder.setSlop(3);
            builder.add(new Term("content", "pingpeng"));
            //第一个Term
            builder.add(new Term("content", "i"));
            //产生距离之后的第二个Term
            PhraseQuery query = builder.build();
            TopDocs tds = searcher.search(query, num);
            System.out.println("一共查询了:" + tds.totalHits);
            for (ScoreDoc sd : tds.scoreDocs) {
                Document doc = searcher.doc(sd.doc);
                System.out.println(doc.get("id") + "---->" +
                        doc.get("name") + "[" + doc.get("email") + "]-->" + doc.get("id") + "," +
                        doc.get("attach") + "," + doc.get("date"));
            }
            this.reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 容错查询
     *
     * @param num
     */
    public void searchByFuzzy(int num) {
        try {
            initIndexReader();
            IndexSearcher searcher = new IndexSearcher(this.reader);
            FuzzyQuery query = new FuzzyQuery(new Term("name", "mase"), 1, 0);
            System.out.println(query.getPrefixLength());
            TopDocs tds = searcher.search(query, num);
            System.out.println("一共查询了:" + tds.totalHits);
            for (ScoreDoc sd : tds.scoreDocs) {
                Document doc = searcher.doc(sd.doc);
                System.out.println(doc.get("id") + "---->" +
                        doc.get("name") + "[" + doc.get("email") + "]-->" + doc.get("id") + "," +
                        doc.get("attach") + "," + doc.get("date"));
            }
            this.reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * QueryParse查询,实际使用比较多
     *
     * @param query
     * @param num
     */
    public void searchByQueryParse(Query query, int num) {
        try {
            initIndexReader();
            IndexSearcher searcher = new IndexSearcher(this.reader);
            TopDocs tds = searcher.search(query, num);
            System.out.println("一共查询了:" + tds.totalHits);
            for (ScoreDoc sd : tds.scoreDocs) {
                Document doc = searcher.doc(sd.doc);
                System.out.println(doc.get("id") + "---->" +
                        doc.get("name") + "[" + doc.get("email") + "]-->" + doc.get("id") + "," +
                        doc.get("attach") + "," + doc.get("date") + "==" + sd.score);
            }
            this.reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 分页查询,方法一
     *
     * @param query
     * @param pageIndex
     * @param pageSize
     */
    public void searchPage(String query, int pageIndex, int pageSize) {
        try {
            initIndexReader();
            IndexSearcher searcher = new IndexSearcher(this.reader);
            QueryParser parser = new QueryParser("content", new StandardAnalyzer());
            Query q = parser.parse(query);
            TopDocs tds = searcher.search(q, 500);
            ScoreDoc[] sds = tds.scoreDocs;
            int start = (pageIndex - 1) * pageSize;
            int end = pageIndex * pageSize;
            for (int i = start; i < end; i++) {
                Document doc = searcher.doc(sds[i].doc);
                System.out.println(sds[i].doc + ":" + doc.get("path") + "-->" + doc.get("filename"));
            }

            this.reader.close();
        } catch (IOException | ParseException e) {
            e.printStackTrace();
        }
    }

    /**
     * 根据页码和分页大小获取上一次的最后一个ScoreDoc
     */
    private ScoreDoc getLastScoreDoc(int pageIndex, int pageSize, Query query, IndexSearcher searcher) throws IOException {
        if (pageIndex == 1) return null;//如果是第一页就返回空
        int num = pageSize * (pageIndex - 1);//获取上一页的数量
        TopDocs tds = searcher.search(query, num);
        return tds.scoreDocs[num - 1];
    }

    /**
     * 分页查询,方法二
     *
     * @param query
     * @param pageIndex
     * @param pageSize
     */
    public void searchPageByAfter(String query, int pageIndex, int pageSize) {
        try {
            initIndexReader();
            IndexSearcher searcher = new IndexSearcher(this.reader);
            QueryParser parser = new QueryParser("content", new StandardAnalyzer());
            Query q = parser.parse(query);
            //先获取上一页的最后一个元素
            ScoreDoc lastSd = getLastScoreDoc(pageIndex, pageSize, q, searcher);
            //通过最后一个元素搜索下页的pageSize个元素
            TopDocs tds = searcher.searchAfter(lastSd, q, pageSize);
            for (ScoreDoc sd : tds.scoreDocs) {
                Document doc = searcher.doc(sd.doc);
                System.out.println(sd.doc + ":" + doc.get("path") + "-->" + doc.get("filename"));
            }
            this.reader.close();
        } catch (ParseException | IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 不分页查询
     *
     * @param query
     */
    public void searchNoPage(String query) {
        try {
            initIndexReader();
            IndexSearcher searcher = new IndexSearcher(this.reader);
            QueryParser parser = new QueryParser("content", new StandardAnalyzer());
            Query q = parser.parse(query);
            TopDocs tds = searcher.search(q, 20);
            ScoreDoc[] sds = tds.scoreDocs;
            for (int i = 0; i < sds.length; i++) {
                Document doc = searcher.doc(sds[i].doc);
                System.out.println(sds[i].doc + ":" + doc.get("path") + "-->" + doc.get("filename"));
            }

            this.reader.close();
        } catch (ParseException | IOException e) {
            e.printStackTrace();
        }
    }
}
